# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""export"""
import argparse
import os

import numpy as np
from mindspore import Tensor
from mindspore import context
from mindspore import dtype as mstype
from mindspore import load_checkpoint
from mindspore import load_param_into_net
from mindspore.train.serialization import export
import sys
cpath = os.getcwd()
sys.path.append(os.getcwd())
from mind3d.dataset.kitti.src.utils import get_config
from mind3d.dataset.kitti.src.utils import get_model_dataset
from mind3d.dataset.kitti.src.utils import get_params_for_net
from mind3d.utils.load_yaml import load_yaml


def run_export(opt):   
    if opt['model_name'] == 'groupfree_3d':
        from mind3d.models.groupfree_3d import Groupfree3DModel
        from mind3d.dataset.scannet_v2 import ScannetDatasetConfig
        parser = argparse.ArgumentParser(description='MindSpore Pointnet')
        parser.add_argument('--device_target', default='GPU', help='device')
        # parser.add_argument('--model', type=str, default='ckpt/best.ckpt', help='model path')
        parser.add_argument('--file_format', type=str, default='MINDIR', help="export file format")
        parser.add_argument('-opt', type=str, default='config/groupfree_3d/scannet.yaml')
        args = parser.parse_known_args()[0]
        opt = load_yaml(args.opt)
        args = parser.parse_args()

        DATASET_CONFIG = ScannetDatasetConfig()


        context.set_context(mode=context.PYNATIVE_MODE, device_target=args.device_target)

        if not os.path.exists('./mindir'):
            os.mkdir('./mindir')

        """select model"""
        if opt['model'] == 'groupfree_3d':
            num_input_channel = 0
            detector = Groupfree3DModel(num_class=DATASET_CONFIG.num_class,
                                        num_heading_bin=DATASET_CONFIG.num_heading_bin,
                                        num_size_cluster=DATASET_CONFIG.num_size_cluster,
                                        mean_size_arr=DATASET_CONFIG.mean_size_arr,
                                        input_feature_dim=num_input_channel,
                                        width=opt['width'],
                                        bn_momentum=opt['bn_momentum'],
                                        sync_bn=False,
                                        num_proposal=opt['num_target'],
                                        sampling=opt['sampling'],
                                        dropout=opt['transformer_dropout'],
                                        activation=opt['transformer_activation'],
                                        nhead=opt['nhead'],
                                        num_decoder_layers=opt['num_decoder_layers'],
                                        dim_feedforward=opt['dim_feedforward'],
                                        self_position_embedding=opt['self_position_embedding'],
                                        cross_position_embedding=opt['cross_position_embedding'],
                                        size_cls_agnostic=False)
        load_checkpoint(opt['checkpoint_path'], net=detector)
        # input_data = Tensor(np.random.uniform(0.0, 1.0, size=[2, 20000, 3]), ms.float32)
        if opt['model'] == 'groupfree_3d':
            input_data = Tensor(np.ones((2, 20000, 3)), ms.float32)
        else:
            pass
        # input_data = ms.Tensor(shape=[None, 20000, 3], dtype=ms.float32)
        export(detector, input_data, file_name='./mindir/detection', file_format=args.file_format)
        print("successfully export model")
    elif opt['model_name'] == 'pointpilliars':
        cfg_path = opt['export_set_up']['cfg_path'] 
        ckpt_path = opt['export_set_up']['ckpt_path'] 
        file_name = opt['export_set_up']['file_name'] 
        file_format = opt['export_set_up']['file_format'] 
        GPU_ID = opt['export_set_up']['GPU_ID'] 

        cfg = get_config(cfg_path)

        device_target = cfg['train_config']['device_target']
        device_id = int(os.getenv('DEVICE_ID', GPU_ID))
        context.set_context(mode=context.GRAPH_MODE, device_target=device_target, device_id=device_id)

        pointpillarsnet, _ = get_model_dataset(cfg)

        params = load_checkpoint(ckpt_path)
        new_params = get_params_for_net(params)
        load_param_into_net(pointpillarsnet, new_params)

        v = cfg['eval_input_reader']['max_number_of_voxels']
        p = cfg['model']['voxel_generator']['max_number_of_points_per_voxel']
        n = cfg['model']['num_point_features']
        
        voxels = Tensor(np.zeros((1, v, p, n)), mstype.float32)
        num_points = Tensor(np.zeros((1, v)), mstype.int32)
        coors = Tensor(np.zeros((1, v, 4)), mstype.int32)
        
        if not os.path.exists('./export_dir'):
            os.mkdir('./export_dir')
        
        if cfg['model']['use_bev']:
            pc_range = np.array(cfg['model']['voxel_generator']['point_cloud_range'])
            voxel_size = np.array(cfg['model']['voxel_generator']['voxel_size'])
            x, y, z = ((pc_range[3:] - pc_range[:3]) / voxel_size).astype('int32')
            bev_map = Tensor(np.zeros((1,) + (z, x * 2, y * 2)), mstype.float32)

            export(pointpillarsnet, voxels, num_points, coors, bev_map, file_name=file_name, file_format=file_format)
        else:
            export(pointpillarsnet, voxels, num_points, coors, file_name=file_name, file_format=file_format)

        print('PointPillars model exported successfully!', 'saved as:',file_name)
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='3d detection export')
    parser.add_argument('--opt', default='/data1/mind3d-master/configs/pointpillars/car_xyres16.yaml', help='Path to config file.')
    args = parser.parse_known_args()[0]
    opt = load_yaml(args.opt)
    run_export(opt)
